#include "CacheManager.h"
#include <limits>
#include <cmath>
#include <algorithm>
#include <iostream>
#include <iomanip>

CacheManager::CacheManager(long long capacity) : capacity_(capacity) {
    if (capacity > 0) {
        free_blocks_.push_back({ 0, capacity, -1 });
    }
}

long long CacheManager::allocate(long long size_needed, int buf_id, const std::string& strategy,
    const std::unordered_map<int, int>& buf_id_to_free_step,
    int current_buf_free_step) {
    if (size_needed <= 0) return -1;

    int best_block_idx = -1;

    if (strategy == "BestFit") {
        long long min_suitable_size = std::numeric_limits<long long>::max();
        for (int i = 0; i < free_blocks_.size(); ++i) {
            if (free_blocks_[i].size >= size_needed && free_blocks_[i].size < min_suitable_size) {
                min_suitable_size = free_blocks_[i].size;
                best_block_idx = i;
            }
        }
    }
    else if (strategy == "LowAddress") {
        long long min_address = std::numeric_limits<long long>::max();
        for (int i = 0; i < free_blocks_.size(); ++i) {
            if (free_blocks_[i].size >= size_needed) {
                if (free_blocks_[i].start < min_address) {
                    min_address = free_blocks_[i].start;
                    best_block_idx = i;
                }
            }
        }
    }
    else if (strategy == "LifetimeAware") {
        double best_score = std::numeric_limits<double>::infinity();
        std::vector<int> suitable_indices;
        for (int i = 0; i < free_blocks_.size(); ++i) {
            if (free_blocks_[i].size >= size_needed) {
                suitable_indices.push_back(i);
            }
        }

        if (!suitable_indices.empty()) {
            if (allocated_blocks_.empty()) { // Fallback to BestFit
                long long min_suitable_size = std::numeric_limits<long long>::max();
                for (int idx : suitable_indices) {
                    if (free_blocks_[idx].size < min_suitable_size) {
                        min_suitable_size = free_blocks_[idx].size;
                        best_block_idx = idx;
                    }
                }
            }
            else {
                for (int idx : suitable_indices) {
                    const auto& fb = free_blocks_[idx];
                    int pred_neighbor_freetime = current_buf_free_step;
                    int succ_neighbor_freetime = current_buf_free_step;

                    // Find predecessor neighbor
                    for (int j = allocated_blocks_.size() - 1; j >= 0; --j) {
                        if (allocated_blocks_[j].start + allocated_blocks_[j].size == fb.start) {
                            if (buf_id_to_free_step.count(allocated_blocks_[j].buf_id)) {
                                pred_neighbor_freetime = buf_id_to_free_step.at(allocated_blocks_[j].buf_id);
                            }
                            break;
                        }
                    }
                    // Find successor neighbor
                    for (const auto& ab : allocated_blocks_) {
                        if (ab.start == fb.start + fb.size) {
                            if (buf_id_to_free_step.count(ab.buf_id)) {
                                succ_neighbor_freetime = buf_id_to_free_step.at(ab.buf_id);
                            }
                            break;
                        }
                    }
                    double score = std::abs(static_cast<double>(current_buf_free_step) - pred_neighbor_freetime) +
                        std::abs(static_cast<double>(current_buf_free_step) - succ_neighbor_freetime);
                    if (score < best_score) {
                        best_score = score;
                        best_block_idx = idx;
                    }
                }
            }
        }
    }

    if (best_block_idx == -1) return -1;

    MemBlock best_block = free_blocks_[best_block_idx];
    long long offset = best_block.start;
    allocated_blocks_.push_back({ offset, size_needed, buf_id });
    std::sort(allocated_blocks_.begin(), allocated_blocks_.end());

    if (best_block.size > size_needed) {
        free_blocks_[best_block_idx].start += size_needed;
        free_blocks_[best_block_idx].size -= size_needed;
    }
    else {
        free_blocks_.erase(free_blocks_.begin() + best_block_idx);
    }

    return offset;
}

bool CacheManager::allocateAt(long long size_needed, int buf_id, long long requested_offset) {
    if (size_needed <= 0 || (requested_offset + size_needed) > capacity_) {
        return false;
    }

    for (int i = 0; i < free_blocks_.size(); ++i) {
        const auto& fb_ref = free_blocks_[i];
        if (requested_offset >= fb_ref.start && (requested_offset + size_needed) <= (fb_ref.start + fb_ref.size)) {

            // *** MODIFIED: Make a copy of the block BEFORE erasing to avoid using an invalidated reference ***
            MemBlock block_to_split = fb_ref;

            free_blocks_.erase(free_blocks_.begin() + i);

            // *** MODIFIED: Use the safe copy 'block_to_split' for all subsequent calculations ***
            if (block_to_split.start < requested_offset) {
                free_blocks_.push_back({ block_to_split.start, requested_offset - block_to_split.start, -1 });
            }

            if ((block_to_split.start + block_to_split.size) > (requested_offset + size_needed)) {
                free_blocks_.push_back({ requested_offset + size_needed,
                                        (block_to_split.start + block_to_split.size) - (requested_offset + size_needed), -1 });
            }

            allocated_blocks_.push_back({ requested_offset, size_needed, buf_id });

            std::sort(free_blocks_.begin(), free_blocks_.end());
            std::sort(allocated_blocks_.begin(), allocated_blocks_.end());

            return true;
        }
    }
    return false;
}

void CacheManager::free(int buf_id_to_free) {
    auto it = std::find_if(allocated_blocks_.begin(), allocated_blocks_.end(),
        [buf_id_to_free](const MemBlock& b) { return b.buf_id == buf_id_to_free; });

    if (it == allocated_blocks_.end()) return;

    MemBlock new_free_block = *it;
    allocated_blocks_.erase(it);

    free_blocks_.push_back(new_free_block);
    std::sort(free_blocks_.begin(), free_blocks_.end());

    mergeFreeBlocks();
}

void CacheManager::mergeFreeBlocks() {
    if (free_blocks_.size() < 2) return;

    std::vector<MemBlock> new_free_list;
    new_free_list.push_back(free_blocks_[0]);

    for (size_t i = 1; i < free_blocks_.size(); ++i) {
        if (new_free_list.back().start + new_free_list.back().size == free_blocks_[i].start) {
            new_free_list.back().size += free_blocks_[i].size;
        }
        else {
            new_free_list.push_back(free_blocks_[i]);
        }
    }
    free_blocks_ = new_free_list;
}

void CacheManager::printState(const std::string& context, int step, int node_id, int buf_id, const std::string& timing) const {
    std::cout << "\n--- [" << context << " | Step: " << step
        << " | Node: " << node_id << " | BufID: " << buf_id
        << "] --- (" << timing << ")\n";

    std::cout << "  Allocated Blocks (" << allocated_blocks_.size() << "):\n";
    if (allocated_blocks_.empty()) {
        std::cout << "    (None)\n";
    }
    else {
        for (const auto& blk : allocated_blocks_) {
            std::cout << "    - BufID: " << std::left << std::setw(5) << blk.buf_id
                << " | Start: " << std::left << std::setw(5) << blk.start
                << " | Size: " << blk.size << "\n";
        }
    }

    std::cout << "  Free Blocks (" << free_blocks_.size() << "):\n";
    if (free_blocks_.empty()) {
        std::cout << "    (None)\n";
    }
    else {
        for (const auto& blk : free_blocks_) {
            std::cout << "    - Start: " << std::left << std::setw(5) << blk.start
                << " | Size: " << blk.size << "\n";
        }
    }
    std::cout << "-----------------------------------------------------------------------------------\n" << std::flush;
}